from pynvml import *
import time
# 探测GPU温度，一旦过高则进行sleep，一直到温度降到可接受程度
nvmlInit()     #初始化
print("Driver: ", nvmlSystemGetDriverVersion())  #显示驱动信息

deviceCount = nvmlDeviceGetCount()
for i in range(deviceCount):
    handle = nvmlDeviceGetHandleByIndex(i)
    print("GPU", i, ":", nvmlDeviceGetName(handle))

handle = nvmlDeviceGetHandleByIndex(0)
info = nvmlDeviceGetMemoryInfo(handle)
print("Memory Total: ",info.total)
print("Memory Free: ",info.free)
print("Memory Used: ",info.used)

print("Temperature is: {} C".format(nvmlDeviceGetTemperature(handle,0)))
print("Power ststus",nvmlDeviceGetPowerState(handle))
TEMP_TOP=92
TEMP_SAFE=45
temperature = nvmlDeviceGetTemperature(handle,0)
while (temperature >= TEMP_TOP):
    print("+++++++ current temperature is: {}, it's too hot, I need a rest...+++++++".format(temperature))
    time.sleep(3)
    temperature = nvmlDeviceGetTemperature(handle,0)

nvmlShutdown()